package com.btl.service.ai.support.input.common.split.splitters;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.segment.TextSegment;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.StringUtils;

import java.util.List;
import java.util.Optional;


@Slf4j
public class TitledRecursiveSplitter implements DocumentSplitter {

    private final int maxTokens;
    private final int overlap;
    private final String titleKey;

    public TitledRecursiveSplitter(int maxTokens, int overlap, String titleKey) {
        this.maxTokens = maxTokens;
        this.overlap = overlap;
        this.titleKey = titleKey;
    }

    @Override
    public List<TextSegment> split(Document document) {
        List<TextSegment> segments = DocumentSplitters.recursive(maxTokens, overlap)
                .split(document);
        String title = Optional.ofNullable(document.metadata()
                        .getString(titleKey))
                .orElse("");
        if (StringUtils.hasText(title)) {
            return segments.stream()
                    .map(segment -> {
                        String titledText = title + "\n\n" + segment.text();
                        Metadata newMetadata = segment.metadata()
                                .copy();
                        newMetadata.put(titleKey, title);
                        return TextSegment.from(titledText, newMetadata);
                    })
                    .toList();
        }
        log.warn("title is empty: {}", title);
        return segments;
    }

}
